import cv2
import numpy as np
import glob
from tqdm import tqdm
from math import degrees as dg

# 定义棋盘大小: 注意此处是内部的行、列角点个数，不包含最外边两列，否则会出错
chessboard_size = (7, 7)

# 生成54×3的矩阵，用来保存棋盘图中9*6(54)个内角点的3D坐标，也就是物体点坐标
objp = np.zeros((np.prod(chessboard_size), 3), dtype=np.float32)
# 通过np.mgrid生成对象的(x,y)坐标点，每个棋盘格大小是130mm
# 最终得到z=0的objp为(0,0,0), (1*13,0,0), (2*13,0,0) ,...
objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2) * 13

# print("object is %f", objp)

# 定义数组，来保存监测到的点
obj_points = []  # 保存世界坐标系的三维点
img_points = []  # 保存图片坐标系的二维点

# 设置终止条件： 迭代30次或者变动 < 0.001
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# 读取目录下的所有图片
calibration_paths = glob.glob(r'C:/Users/HuangSX/Desktop/OpenCV/data/mytest*.jpg')

# 为方便显示，使用tqdm显示进度条
for image_path in tqdm(calibration_paths):
    # 读取图片
    img = cv2.imread(image_path)

    # 图像二值化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 找到棋盘格内角点位置
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

    if ret:
        obj_points.append(objp)
        # 亚像素级角点检测，在角点检测中精确化角点位置
        corners2 = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)
        img_points.append(corners2)

        # 在图中标注角点,方便查看结果
        img = cv2.drawChessboardCorners(img, chessboard_size, corners2, ret)
        img = cv2.resize(img, (400, 600))
        cv2.imshow('img', img)
        cv2.waitKey(0)

cv2.destroyAllWindows()
print("finish all the pic count")

# 相机标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, gray.shape, None, None)

# 显示和保存参数
print("#######相机内参#######")
print(mtx)
print("#######畸变系数#######")
print(dist)
print("#######相机旋转矩阵#######")
print(rvecs)
print("#######相机平移矩阵#######")
print(tvecs)


# 使用一张图片看看去畸变之后的效果
img2 = cv2.imread(r'C:\Users\HuangSX\Desktop\OpenCV\data\mytest01.jpg')
print("orgininal img_point array shape", img.shape)
# img2.shape[:2]取图片 高、宽；
h, w = img2.shape[:2]
print("pic's hight, weight: %f,  %f" % (h, w))
# img2.shape[:3]取图片的 高、宽、通道
# h,  w ,n= img2.shape[:3]
# print("PIC shape", (h, w, n))

# 自由比例参数
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
dst = cv2.undistort(img2, mtx, dist, None, newcameramtx)

# 根据前面ROI区域裁剪图片
x, y, w, h = roi
dst = dst[y:y + h, x:x + w]
cv2.imwrite('calibresult.jpg', dst)

# 计算所有图片的平均重投影误差
total_error = 0
for i in range(len(obj_points)):
    img_points2, _ = cv2.projectPoints(obj_points[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv2.norm(img_points[i], img_points2, cv2.NORM_L2)/len(img_points2)
    total_error += error
print("total error: {}".format(total_error/len(obj_points)))

"""
使用同一相机，将棋盘格放在前方1m左右固定，然后使用线性方法进行相对位姿估计，
然后评价结果的合理性
"""
# 加载相机标定的内参数、畸变参数矩阵，并打印结果
with np.load('C.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

print(mtx, '\n' * 2, dist)

# 像素坐标
# test_img = cv2.imread(r'C:/Users/HuangSX/Desktop/OpenCV/data/left01.jpg')
test_img = cv2.imread(r'C:/Users/HuangSX/Desktop/OpenCV/data/mytest.jpg')
# 图片转换为灰度图
gray = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)

# 找到图像平面点角点坐标
ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

if ret:
    # 使用Opencv2提供的相机位姿估计接口
    _, R, T, _, = cv2.solvePnPRansac(objp, corners, mtx, dist)
    print("旋转向量", R)
    print("平移向量", T)

sita_x = dg(R[0][0])
sita_y = dg(R[1][0])
sita_z = dg(R[2][0])

print("sita_x is  ", sita_x)
print("sita_y is  ", sita_y)
print("sita_z is  ", sita_z)

